from OCC.Core.BRep import BRep_Builder
from OCC.Core.TopoDS import TopoDS_Compound
from OCC.Core.BRepPrimAPI import BRepPrimAPI_MakeSphere
from OCC.Core.gp import gp_Pnt
from OCC.Core.STEPControl import STEPControl_Writer, STEPControl_StepModelType
from OCC.Core.IFSelect import IFSelect_RetDone
from OCC.Display.SimpleGui import init_display
# 定义球的半径
radius = 10.0

# 定义球心的位置，使得它们在空间中排列成一个立方体的八个顶点
centers = [
    gp_Pnt(-radius, -radius, -radius),
    gp_Pnt(radius, -radius, -radius),
    gp_Pnt(-radius, radius, -radius),
    gp_Pnt(radius, radius, -radius),
    gp_Pnt(-radius, -radius, radius),
    gp_Pnt(radius, -radius, radius),
    gp_Pnt(-radius, radius, radius),
    gp_Pnt(radius, radius, radius)
]

# 创建八个球
spheres = [BRepPrimAPI_MakeSphere(center, radius).Shape() for center in centers]

# 创建一个复合形状来包含所有球
builder = BRep_Builder()
compound = TopoDS_Compound()
builder.MakeCompound(compound)

for sphere in spheres:
    builder.Add(compound, sphere)

# 导出为 STEP 文件
step_writer = STEPControl_Writer()
status = step_writer.Transfer(compound, STEPControl_StepModelType.STEPControl_AsIs)
if status == IFSelect_RetDone:
    step_writer.Write("eight_connected_spheres.stp")
    print("文件已成功导出为 eight_connected_spheres.stp")
else:
    print("导出失败")

display, start_display, add_menu, add_function_to_menu = init_display()

# 将所有球体添加到显示窗口
for sphere in spheres:
    display.DisplayShape(sphere, update=True)

# 启动事件循环以保持窗口打开
start_display()